# coding:utf-8
import os
from PyQt5 import QtCore,QtGui,QtWidgets,Qt
from PyQt5.QtWidgets import QMessageBox
import sys
import qtawesome
import numpy as np
import cv2
from models import *
import models.regnet
# import simpleaudio as sa
import threading
import inspect
import ctypes
from torchvision import transforms
from train import Classification
from PIL import Image

os.environ["CUDA_VISIBLE_DEVICES"] = "0"
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"

class MainUi(QtWidgets.QMainWindow):

    # 窗口最大化信号
    windowMaximumed = QtCore.pyqtSignal()
    # 窗口还原信号
    windowNormaled = QtCore.pyqtSignal()
    # 窗口关闭信号
    windowClosed = QtCore.pyqtSignal()

    def __init__(self):
        super().__init__()
        # self.
        self.clspre1 = Classification(model_name='./history/VGG0.01/model300.pth', train_net=VGG('VGG16'))
        self.clspre2 = Classification(model_name='./history/senet0.01/model300.pth', train_net=SENet18())
        self.clspre3 = Classification(model_name='./history/regnet0.01/model300.pth', train_net=models.regnet.RegNetX_400MF())
        splash = QtWidgets.QSplashScreen(QtGui.QPixmap("./input/image/Icon.png"))
        splash.show() 
        self.init_ui()
        self.bufferImagePath = './input/image/buffer.jpg'
        # self.loadStyleModel()
        splash.finish(self)

    def init_ui(self):
        self.setFixedSize(1000,700)
        self.main_widget = QtWidgets.QWidget() # 创建窗口主部件
        self.main_layout = QtWidgets.QGridLayout() # 创建主部件的网格布局
        self.main_widget.setLayout(self.main_layout) # 设置窗口主部件布局为网格布局

        self.left_widget = QtWidgets.QWidget() # 创建左侧部件
        self.left_widget.setObjectName('left_widget')
        self.left_layout = QtWidgets.QVBoxLayout()
        self.left_widget.setLayout(self.left_layout) # 设置左侧部件布局为网格

        self.right_widget = QtWidgets.QWidget() # 创建右侧部件
        self.right_widget.setObjectName('right_widget')
        self.right_layout = QtWidgets.QGridLayout()
        self.right_widget.setLayout(self.right_layout) # 设置右侧部件布局为网格

        self.main_layout.addWidget(self.left_widget,0,0,1,1) # 左侧部件在第0行第0列，占8行2列
        self.main_layout.addWidget(self.right_widget,0,2,1,10) # 右侧部件在第0行第3列，占8行10列
        self.setCentralWidget(self.main_widget) # 设置窗口主部件

        self.left_top_layout = QtWidgets.QHBoxLayout()
        self.left_close = QtWidgets.QPushButton("") # 关闭按钮
        self.left_close.clicked.connect(self.close)
        self.left_visit = QtWidgets.QPushButton("") # 最大化按钮
        self.left_visit.clicked.connect(self.showMaximized)
        self.left_mini = QtWidgets.QPushButton("") # 最小化按钮
        self.left_mini.clicked.connect(self.showMinimized)
        self.left_top_layout.addWidget(self.left_mini)
        self.left_top_layout.addWidget(self.left_visit)
        self.left_top_layout.addWidget(self.left_close)

        self.left_label_1 = QtWidgets.QPushButton("图像")
        self.left_label_1.setObjectName('left_label')
        self.left_label_2 = QtWidgets.QPushButton("模型")
        self.left_label_2.setObjectName('left_label')
        self.left_label_3 = QtWidgets.QPushButton("存储")
        self.left_label_3.setObjectName('left_label')

        self.left_button_1 = QtWidgets.QPushButton(qtawesome.icon('fa.image',color='white'),"选择图片")
        self.left_button_1.setObjectName('left_button')
        self.left_button_1.clicked.connect(self.selectImage)

        self.left_button_3 = QtWidgets.QPushButton(qtawesome.icon('fa.sellsy',color='white'),"New_VGG")
        self.left_button_3.setObjectName('left_button')
        self.left_button_3.clicked.connect(self.style)
        self.left_button_4 = QtWidgets.QPushButton(qtawesome.icon('fa.image',color='white'),"Inception")
        self.left_button_4.setObjectName('left_button')
        self.left_button_4.clicked.connect(self.styleHaoyao)
        self.left_button_5 = QtWidgets.QPushButton(qtawesome.icon('fa.image',color='white'),"SE-Res-Inception")
        self.left_button_5.setObjectName('left_button')
        self.left_button_5.clicked.connect(self.styleShinkai)
        self.left_button_6 = QtWidgets.QPushButton(qtawesome.icon('fa.download',color='white'),"保存")
        self.left_button_6.setObjectName('left_button')
        self.left_button_6.clicked.connect(self.save)

        self.left_layout.addLayout(self.left_top_layout)
        self.left_layout.addWidget(self.left_label_1)
        self.left_layout.addWidget(self.left_button_1)
        self.left_layout.addWidget(self.left_label_2)
        self.left_layout.addWidget(self.left_button_4)
        self.left_layout.addWidget(self.left_button_5)
        self.left_layout.addWidget(self.left_button_3)
        self.left_layout.addWidget(self.left_label_3)
        self.left_layout.addWidget(self.left_button_6)

        self.mode = 0
        self.right_top_widget = QtWidgets.QWidget()
        # self.right_top_layout = QtWidgets.QHBoxLayout()
        self.right_top_layout = QtWidgets.QGridLayout()
        self.right_top_widget.setLayout(self.right_top_layout)
        self.show_image_label = QtWidgets.QLabel()
        # self.show_image_label.setStyleSheet('''QLabel{background:#000000}''')
        self.show_style_label = QtWidgets.QLabel()
        # self.show_style_label.setStyleSheet('''QLabel{background:#F16F31}''')
        self.show_style_label.setVisible(False)
        self.right_top_layout.addWidget(self.show_image_label, 0, 0, 15, 15)
        self.right_top_layout.addWidget(self.show_style_label, 0, 10, 5, 5)

        self.image = QtGui.QPixmap()
        self.show_image_label.setAlignment(QtCore.Qt.AlignCenter)
        # self.show_image_label.setPixmap(self.image)

        self.style = QtGui.QPixmap()
        self.show_style_label.setAlignment(QtCore.Qt.AlignCenter)
        # self.show_style_label.setPixmap(self.style)
        self.setMainView()

        self.right_bottom_widget = QtWidgets.QWidget()
        self.right_bottom_layout = QtWidgets.QHBoxLayout()
        self.right_bottom_widget.setLayout(self.right_bottom_layout)
        self.right_process_bar = QtWidgets.QSlider()
        self.right_process_bar.setOrientation(QtCore.Qt.Horizontal)
        self.right_process_bar.setValue(0)
        self.right_process_bar.valueChanged.connect(self.s1_changed)
        self.right_process_bar.sliderPressed.connect(self.s1_press)
        self.right_process_bar.sliderReleased.connect(self.s1_release)

        self.console_button_1 = QtWidgets.QPushButton(qtawesome.icon('fa.backward', color='#F76677'), "")
        self.console_button_2 = QtWidgets.QPushButton(qtawesome.icon('fa.pause', color='#F76677'), "")
        self.console_button_3 = QtWidgets.QPushButton(qtawesome.icon('fa.forward', color='#F76677'), "")

        self.console_button_1.setFixedSize(20, 20)
        self.console_button_2.setFixedSize(30, 30)
        self.console_button_3.setFixedSize(20, 20)

        self.console_button_1.setStyleSheet('''QPushButton{border-radius:10px;}QPushButton:hover{border:1px solid yellow;}''')
        self.console_button_2.setStyleSheet('''QPushButton{border-radius:15px;}QPushButton:hover{border:1px solid yellow;}''')
        self.console_button_3.setStyleSheet('''QPushButton{border-radius:10px;}QPushButton:hover{border:1px solid yellow;}''')

        self.right_bottom_layout.addWidget(self.console_button_1)
        self.right_bottom_layout.addWidget(self.console_button_2)
        self.right_bottom_layout.addWidget(self.console_button_3)
        self.right_bottom_layout.addWidget(self.right_process_bar)
        self.right_bottom_widget.setVisible(False)

        self.right_layout.addWidget(self.right_top_widget, 0, 0, 15, 8)
        self.right_layout.addWidget(self.right_bottom_widget, 16, 0, 1, 8)
        
        self.left_close.setFixedSize(20, 20) # 设置关闭按钮的大小
        self.left_visit.setFixedSize(20, 20) # 设置按钮大小
        self.left_mini.setFixedSize(20, 20) # 设置最小化按钮大小

        self.left_close.setStyleSheet('''QPushButton{background:#F76677;border-radius:10px;}QPushButton:hover{background:red;}''')
        self.left_visit.setStyleSheet('''QPushButton{background:#F7D674;border-radius:10px;}QPushButton:hover{background:yellow;}''')
        self.left_mini.setStyleSheet('''QPushButton{background:#6DDF6D;border-radius:10px;}QPushButton:hover{background:green;}''')

        self.left_widget.setStyleSheet('''
            QPushButton{border:none;color:white;}
            QPushButton#left_label{
                border:none;
                border-bottom:1px solid white;
                font-size:18px;
                font-weight:700;
                font-family: "Helvetica Neue", Helvetica, Arial, sans-serif;
            }
            QPushButton#left_button:hover{border-left:4px solid red;font-weight:700;}
            ''')

        self.right_widget.setStyleSheet('''
            QWidget#right_widget{
                background:gray;
                border-top-right-radius:5px;
                border-bottom-right-radius:5px;
            }
            ''')
        
        self.setWindowOpacity(1) # 设置窗口透明度
        self.setAttribute(QtCore.Qt.WA_TranslucentBackground) # 设置窗口背景透明
        self.setWindowFlag(QtCore.Qt.FramelessWindowHint) # 隐藏边框
        self.main_widget.setStyleSheet('''
            QWidget#left_widget{
            background:gray;
            border-top-left-radius:5px;
            border-bottom-left-radius:5px;
            }
            ''')
        self.main_layout.setSpacing(1)

    def showMaximized(self):
        if not self.isMaximized():
            # 最大化
            self.setWindowState(QtCore.Qt.WindowMaximized)
            self.windowMaximumed.emit()
        else:  # 还原
            self.showNormal()
            self.windowNormaled.emit()
        self.setMainView()

    def mousePressEvent(self, event):
        if event.button()==QtCore.Qt.LeftButton:
            self.m_flag=True
            self.m_Position=event.globalPos()-self.pos() #获取鼠标相对窗口的位置
            event.accept()
            self.setCursor(QtGui.QCursor(QtCore.Qt.OpenHandCursor))  #更改鼠标图标
            
    def mouseMoveEvent(self, QMouseEvent):
        if QtCore.Qt.LeftButton and self.m_flag:  
            self.move(QMouseEvent.globalPos()-self.m_Position)#更改窗口位置
            QMouseEvent.accept()
            
    def mouseReleaseEvent(self, QMouseEvent):
        self.m_flag=False
        self.setCursor(QtGui.QCursor(QtCore.Qt.ArrowCursor))
    
    def mouseDoubleClickEvent(self, event):
        if self.show_style_label.isVisible:
            self.mode = not self.mode
            self.setMainView()

    def setMainView(self):
        if self.mode == 0: # 显示原图像
            self.show_image_label.setPixmap(self.image.scaled(self.show_image_label.size(), aspectRatioMode=QtCore.Qt.KeepAspectRatio, transformMode=QtCore.Qt.SmoothTransformation))
            self.show_style_label.setPixmap(self.style.scaled(self.show_style_label.size(), aspectRatioMode=QtCore.Qt.KeepAspectRatio, transformMode=QtCore.Qt.SmoothTransformation))
        else: # 风格化
            self.show_image_label.setPixmap(self.style.scaled(self.show_image_label.size(), aspectRatioMode=QtCore.Qt.KeepAspectRatio, transformMode=QtCore.Qt.SmoothTransformation))
            self.show_style_label.setPixmap(self.image.scaled(self.show_style_label.size(), aspectRatioMode=QtCore.Qt.KeepAspectRatio, transformMode=QtCore.Qt.SmoothTransformation))

    # def changeImageSize(self):
    #     self.showImage()
    #     self.showStyle()

    # def showImage(self):
    #     self.show_image_label.setPixmap(self.image.scaled(self.show_image_label.size(), aspectRatioMode=QtCore.Qt.KeepAspectRatio, transformMode=QtCore.Qt.SmoothTransformation))

    # def showStyle(self):
    #     self.show_style_label.setPixmap(self.style.scaled(self.show_style_label.size(), aspectRatioMode=QtCore.Qt.KeepAspectRatio, transformMode=QtCore.Qt.SmoothTransformation))

    def selectImage(self):
        self.file_name, file_type = QtWidgets.QFileDialog.getOpenFileName(self,"open file dialog","./verify_image","Image files(*.jpg *.png)")
        self.image = QtGui.QPixmap(self.file_name)
        self.mode = 0
        self.flag = 0
        self.setMainView()
        self.show_style_label.setVisible(False)
        self.right_bottom_widget.setVisible(False)

    def s1_changed(self):
        # print(self.ui.s1.value())
        self.changed = True

    def s1_press(self):
        self.isMousePress = True

    def s1_release(self):
        self.isMousePress = False



    def close(self):
        self.stop_thread(self.thr)
        exit()

    def _async_raise(self, tid, exctype):
        """raises the exception, performs cleanup if needed"""
        tid = ctypes.c_long(tid)
        if not inspect.isclass(exctype):
            exctype = type(exctype)
        res = ctypes.pythonapi.PyThreadState_SetAsyncExc(tid, ctypes.py_object(exctype))
        if res == 0:
            raise ValueError("invalid thread id")
        elif res != 1:
            # """if it returns a number greater than one, you're in trouble,
            # and you should call it again with exc=NULL to revert the effect"""
            ctypes.pythonapi.PyThreadState_SetAsyncExc(tid, None)
            raise SystemError("PyThreadState_SetAsyncExc failed")

    def stop_thread(self, thread):
        self._async_raise(thread.ident, SystemExit)

    
    def style(self):

        transform = transforms.Compose([transforms.Resize((32, 32)), transforms.ToTensor()])
        # 测试图片
        img = Image.open(self.file_name)
        result = self.clspre1.predict(img, transform).cpu().numpy()
        if (result[0] == 1):
            QMessageBox.about(self, "Result", "瑞玉")
        elif (result[0] == 2):
            QMessageBox.about(self, "Result", "徐香")
        else:
            QMessageBox.about(self, "Result", "翠香")
    
    def styleHaoyao(self):

        transform = transforms.Compose([transforms.Resize((32, 32)), transforms.ToTensor()])
        # 测试图片
        img = Image.open(self.file_name)
        result = self.clspre2.predict(img, transform).cpu().numpy()
        if (result[0] == 1):
            QMessageBox.about(self, "Result", "瑞玉")
        elif (result[0] == 2):
            QMessageBox.about(self, "Result", "徐香")
        else:
            QMessageBox.about(self, "Result", "翠香")
    
    def styleShinkai(self):
        if len(self.file_name) == 0:
            QtWidgets.QMessageBox.about(self,"Warning!","Don't Select File...")
            return
        # clspre = Classification(model_name='./history/regnet0.01/model300.pth',train_net=models.regnet.RegNetX_400MF())
        transform = transforms.Compose([transforms.Resize((32, 32)), transforms.ToTensor()])
        # 测试图片
        img = Image.open(self.file_name)
        result = self.clspre3.predict(img, transform).cpu().numpy()
        if (result[0] == 1):
            QMessageBox.about(self, "Result", "瑞玉")
        elif (result[0] == 2):
            QMessageBox.about(self, "Result", "徐香")
        else:
            QMessageBox.about(self, "Result", "翠香")
        
    def save(self):        
        filename, filetype = QtWidgets.QFileDialog.getSaveFileName(None, "Save File Dialog", "./download", "JPG FILES(*.jpg);;PNG FILES(*.png)")
        success = self.style.toImage().save(filename)
        if success == True:
            QtWidgets.QMessageBox.about(self,"Message","Save Successfully!")
        else:
            QtWidgets.QMessageBox.about(self,"Message","Save Failed!")

    def loadStyleModel(self):
        self.sess_h = tf.Session()
        with gfile.FastGFile("./output_model/Hayao_model/frozen_model_50.pb",'rb') as f:
            graph_def = tf.GraphDef()
            graph_def.ParseFromString(f.read())
            self.sess_h.graph.as_default()
            tf.import_graph_def(graph_def,name='')
        self.sess_h.run(tf.global_variables_initializer()) 
        self.input_h = self.sess_h.graph.get_tensor_by_name('test:0')
        self.output_h = self.sess_h.graph.get_tensor_by_name('generator/G_MODEL/Tanh:0')
        self.sess_h.run(self.output_h, feed_dict = {self.input_h : np.asarray(load_test_data('./input/image/runtest.jpg'))})
        
        self.sess_s = tf.Session()
        with gfile.FastGFile("./output_model/Shinkai_model/frozen_model_100.pb",'rb') as f:
            graph_def = tf.GraphDef()
            graph_def.ParseFromString(f.read())
            self.sess_s.graph.as_default()
            tf.import_graph_def(graph_def,name='')
        self.sess_s.run(tf.global_variables_initializer()) 
        self.input_s = self.sess_s.graph.get_tensor_by_name('test:0')
        self.output_s = self.sess_s.graph.get_tensor_by_name('generator/G_MODEL/Tanh:0')
        self.sess_s.run(self.output_s, feed_dict = {self.input_s : np.asarray(load_test_data('./input/image/runtest.jpg'))})
        

def main():
  app = QtWidgets.QApplication(sys.argv)
  app.setWindowIcon(Qt.QIcon("./input/image/icon.ico"))
  gui = MainUi()
  gui.show()
  sys.exit(app.exec_())

if __name__ == '__main__':
  main()